"""
XML处理工具类，提供XML解析、节点操作、内容修改等功能，使用lxml库实现
"""

from lxml import etree
from typing import List, Union, Optional, Dict, Any
from lxml.etree import ElementBase, XMLParser
import os


class XmlUtils:
    """
    XML处理工具类，提供各种XML操作功能，基于lxml实现
    """
    
    def __init__(self, xml_path: str = None, xml_string: str = None):
        """
        初始化XML处理工具

        :param xml_path: XML文件路径，可选
        :param xml_string: XML字符串，可选
        """
        self.tree = None
        self.root = None
        self.namespaces = {}  # 存储命名空间映射
        
        if xml_path:
            self.load_from_file(xml_path)
        elif xml_string:
            self.load_from_string(xml_string)
    
    def load_from_file(self, xml_path: str) -> None:
        """
        从文件加载XML

        :param xml_path: XML文件路径
        """
        if not os.path.exists(xml_path):
            raise FileNotFoundError(f"XML文件不存在: {xml_path}")
        
        # 使用lxml解析XML
        parser = XMLParser(strip_cdata=False)
        self.tree = etree.parse(xml_path, parser)
        self.root = self.tree.getroot()
        self._parse_namespaces()
    
    def load_from_string(self, xml_string: str) -> None:
        """
        从字符串加载XML

        :param xml_string: XML字符串
        """
        parser = XMLParser(strip_cdata=False)
        self.root = etree.fromstring(xml_string.encode('utf-8'), parser)
        self.tree = etree.ElementTree(self.root)
        self._parse_namespaces()
    
    def _parse_namespaces(self) -> None:
        """
        解析XML中的命名空间
        """
        if self.root is not None:
            # 获取所有命名空间
            self.namespaces = self.root.nsmap.copy()
            # 处理默认命名空间
            if None in self.namespaces:
                self.namespaces['default'] = self.namespaces.pop(None)
    
    def save_to_file(self, xml_path: str, encoding: str = 'utf-8', 
                     pretty_print: bool = True) -> None:
        """
        保存XML到文件

        :param xml_path: 保存路径
        :param encoding: 编码格式
        :param pretty_print: 是否美化输出
        """
        if self.tree is None:
            raise ValueError("没有可保存的XML数据")
        
        # 确保目录存在
        directory = os.path.dirname(xml_path)
        if directory and not os.path.exists(directory):
            os.makedirs(directory)
        
        # 保存XML文件
        self.tree.write(xml_path, encoding=encoding, pretty_print=pretty_print, 
                       xml_declaration=True)
    
    def get_root(self) -> ElementBase:
        """
        获取XML根节点

        :return: 根节点
        """
        return self.root
    
    def find_element(self, xpath: str) -> Optional[ElementBase]:
        """
        查找单个元素

        :param xpath: XPath表达式
        :return: 匹配的元素，如果未找到则返回None
        """
        if self.root is None:
            return None
        elements = self.root.xpath(xpath, namespaces=self.namespaces)
        return elements[0] if elements else None
    
    def find_elements(self, xpath: str) -> List[ElementBase]:
        """
        查找多个元素

        :param xpath: XPath表达式
        :return: 匹配的元素列表
        """
        if self.root is None:
            return []
        return self.root.xpath(xpath, namespaces=self.namespaces)
    
    def get_element_text(self, xpath: str, default: str = '') -> str:
        """
        获取元素文本内容

        :param xpath: XPath表达式
        :param default: 默认值
        :return: 元素文本内容
        """
        element = self.find_element(xpath)
        if element is not None and element.text is not None:
            return element.text
        return default
    
    def set_element_text(self, xpath: str, text: str) -> bool:
        """
        设置元素文本内容

        :param xpath: XPath表达式
        :param text: 文本内容
        :return: 是否设置成功
        """
        element = self.find_element(xpath)
        if element is not None:
            element.text = text
            return True
        return False
    
    def get_element_attribute(self, xpath: str, attribute: str, default: str = '') -> str:
        """
        获取元素属性值

        :param xpath: XPath表达式
        :param attribute: 属性名
        :param default: 默认值
        :return: 属性值
        """
        element = self.find_element(xpath)
        if element is not None:
            return element.get(attribute, default)
        return default
    
    def set_element_attribute(self, xpath: str, attribute: str, value: str) -> bool:
        """
        设置元素属性值

        :param xpath: XPath表达式
        :param attribute: 属性名
        :param value: 属性值
        :return: 是否设置成功
        """
        element = self.find_element(xpath)
        if element is not None:
            element.set(attribute, value)
            return True
        return False
    
    def element_exists(self, xpath: str) -> bool:
        """
        检查元素是否存在

        :param xpath: XPath表达式
        :return: 元素是否存在
        """
        return self.find_element(xpath) is not None
    
    def add_element(self, parent_xpath: str, tag: str, 
                    text: str = None, attributes: dict = None) -> Optional[ElementBase]:
        """
        添加子元素

        :param parent_xpath: 父元素XPath
        :param tag: 元素标签名
        :param text: 元素文本内容
        :param attributes: 元素属性字典
        :return: 新添加的元素
        """
        parent = self.find_element(parent_xpath)
        if parent is None:
            # 如果父元素不存在，创建根元素
            if self.root is None:
                self.root = etree.Element(tag)
                self.tree = etree.ElementTree(self.root)
                parent = self.root
            else:
                return None
        else:
            # 创建并添加子元素
            new_element = etree.SubElement(parent, tag)
            if text is not None:
                new_element.text = text
            if attributes is not None:
                for key, value in attributes.items():
                    new_element.set(key, value)
            return new_element
    
    def remove_element(self, xpath: str) -> bool:
        """
        删除元素

        :param xpath: XPath表达式
        :return: 是否删除成功
        """
        if self.root is None:
            return False
        
        # 查找要删除的元素
        elements = self.find_elements(xpath)
        if not elements:
            return False
        
        # 删除所有匹配的元素
        for element in elements:
            # 查找父元素
            parent = element.getparent()
            if parent is not None:
                parent.remove(element)
            elif element == self.root:
                # 如果要删除的是根元素
                self.root = None
                self.tree = None
                return True
        return True
    
    def add_namespace(self, prefix: str, uri: str) -> None:
        """
        添加命名空间

        :param prefix: 命名空间前缀
        :param uri: 命名空间URI
        """
        self.namespaces[prefix] = uri
    
    def register_namespaces(self, namespaces: Dict[str, str]) -> None:
        """
        注册多个命名空间

        :param namespaces: 命名空间字典 {prefix: uri}
        """
        for prefix, uri in namespaces.items():
            self.add_namespace(prefix, uri)
    
    def get_namespaces(self) -> Dict[str, str]:
        """
        获取已注册的命名空间

        :return: 命名空间字典
        """
        return self.namespaces.copy()
    
    def get_all_text(self) -> str:
        """
        获取所有文本内容

        :return: 所有文本内容
        """
        if self.root is None:
            return ""
        
        return self.root.xpath("string()")[0] if self.root.xpath("string()") else ""
    
    def to_string(self, pretty_print: bool = True) -> str:
        """
        将XML转换为字符串

        :param pretty_print: 是否美化输出
        :return: XML字符串
        """
        if self.root is None:
            return ""
        
        return etree.tostring(self.root, pretty_print=pretty_print, encoding='unicode')
    
    def get_element_children(self, xpath: str) -> List[ElementBase]:
        """
        获取元素的子元素列表

        :param xpath: XPath表达式
        :return: 子元素列表
        """
        element = self.find_element(xpath)
        if element is not None:
            return list(element)
        return []
    
    def get_element_count(self, xpath: str) -> int:
        """
        获取匹配元素的数量

        :param xpath: XPath表达式
        :return: 元素数量
        """
        return len(self.find_elements(xpath))
    
    def update_element(self, xpath: str, text: str = None, 
                       attributes: dict = None) -> bool:
        """
        更新元素的文本和属性

        :param xpath: XPath表达式
        :param text: 新的文本内容
        :param attributes: 新的属性字典
        :return: 是否更新成功
        """
        element = self.find_element(xpath)
        if element is None:
            return False
        
        if text is not None:
            element.text = text
        
        if attributes is not None:
            for key, value in attributes.items():
                element.set(key, value)
        
        return True
    
    def copy_element(self, source_xpath: str, target_xpath: str, 
                     new_tag: str = None) -> bool:
        """
        复制元素到指定位置

        :param source_xpath: 源元素XPath
        :param target_xpath: 目标父元素XPath
        :param new_tag: 新元素标签名，如果为None则使用原标签名
        :return: 是否复制成功
        """
        source_elements = self.find_elements(source_xpath)
        target_element = self.find_element(target_xpath)
        
        if not source_elements or target_element is None:
            return False
        
        for source_element in source_elements:
            # 深度复制元素
            copied_element = self._deep_copy_element(source_element, target_element, new_tag)
            target_element.append(copied_element)
        
        return True
    
    def _deep_copy_element(self, source_element: ElementBase, 
                          target_parent: ElementBase, new_tag: str = None) -> ElementBase:
        """
        深度复制元素

        :param source_element: 源元素
        :param target_parent: 目标父元素
        :param new_tag: 新标签名
        :return: 复制的元素
        """
        # 创建新元素
        tag = new_tag if new_tag is not None else source_element.tag
        new_element = etree.SubElement(target_parent, tag)
        
        # 复制文本
        new_element.text = source_element.text
        new_element.tail = source_element.tail
        
        # 复制属性
        for key, value in source_element.attrib.items():
            new_element.set(key, value)
        
        # 递归复制子元素
        for child in source_element:
            self._deep_copy_element(child, new_element)
        
        return new_element
    
    @staticmethod
    def create_empty_xml(root_tag: str = 'root', namespaces: Dict[str, str] = None) -> 'XmlUtils':
        """
        创建空的XML文档

        :param root_tag: 根元素标签名
        :param namespaces: 命名空间字典
        :return: XmlUtils实例
        """
        xml_utils = XmlUtils()
        xml_utils.root = etree.Element(root_tag)
        xml_utils.tree = etree.ElementTree(xml_utils.root)
        
        # 添加命名空间
        if namespaces:
            for prefix, uri in namespaces.items():
                xml_utils.root.nsmap[prefix] = uri
                xml_utils.namespaces[prefix] = uri
        
        return xml_utils
    
    @staticmethod
    def from_dict(data: dict, root_tag: str = 'root') -> 'XmlUtils':
        """
        从字典创建XML

        :param data: 字典数据
        :param root_tag: 根元素标签名
        :return: XmlUtils实例
        """
        def _dict_to_xml(parent, data_dict):
            for key, value in data_dict.items():
                if isinstance(value, dict):
                    element = etree.SubElement(parent, key)
                    _dict_to_xml(element, value)
                elif isinstance(value, list):
                    for item in value:
                        if isinstance(item, dict):
                            element = etree.SubElement(parent, key)
                            _dict_to_xml(element, item)
                        else:
                            element = etree.SubElement(parent, key)
                            element.text = str(item)
                else:
                    element = etree.SubElement(parent, key)
                    element.text = str(value)
        
        xml_utils = XmlUtils()
        xml_utils.root = etree.Element(root_tag)
        xml_utils.tree = etree.ElementTree(xml_utils.root)
        _dict_to_xml(xml_utils.root, data)
        return xml_utils
    
    def create_soap_envelope(self, namespaces: Dict[str, str] = None) -> None:
        """
        创建SOAP信封结构

        :param namespaces: 可选的命名空间字典
        """
        # 默认SOAP命名空间
        default_namespaces = {
            'soap': 'http://schemas.xmlsoap.org/soap/envelope/'
        }
        
        if namespaces:
            default_namespaces.update(namespaces)
        
        # 注册命名空间
        self.register_namespaces(default_namespaces)
        
        # 创建SOAP结构
        self.root = etree.Element(f'{{{default_namespaces["soap"]}}}Envelope', nsmap=default_namespaces)
        self.tree = etree.ElementTree(self.root)
        
        # 创建Header和Body
        header = etree.SubElement(self.root, f'{{{default_namespaces["soap"]}}}Header')
        body = etree.SubElement(self.root, f'{{{default_namespaces["soap"]}}}Body')
    
    def add_soap_header_element(self, tag: str, text: str = None, 
                               attributes: dict = None, namespace: str = None) -> Optional[ElementBase]:
        """
        添加SOAP Header元素

        :param tag: 元素标签名
        :param text: 元素文本内容
        :param attributes: 元素属性字典
        :param namespace: 元素命名空间URI
        :return: 新添加的元素
        """
        header_elements = self.find_elements('//soap:Header')
        if not header_elements:
            return None
        
        header = header_elements[0]
        
        # 处理命名空间
        full_tag = f'{{{namespace}}}{tag}' if namespace else tag
        new_element = etree.SubElement(header, full_tag)
        
        if text is not None:
            new_element.text = text
        if attributes is not None:
            for key, value in attributes.items():
                new_element.set(key, value)
        
        return new_element
    
    def add_soap_body_element(self, tag: str, text: str = None, 
                             attributes: dict = None, namespace: str = None) -> Optional[ElementBase]:
        """
        添加SOAP Body元素

        :param tag: 元素标签名
        :param text: 元素文本内容
        :param attributes: 元素属性字典
        :param namespace: 元素命名空间URI
        :return: 新添加的元素
        """
        body_elements = self.find_elements('//soap:Body')
        if not body_elements:
            return None
        
        body = body_elements[0]
        
        # 处理命名空间
        full_tag = f'{{{namespace}}}{tag}' if namespace else tag
        new_element = etree.SubElement(body, full_tag)
        
        if text is not None:
            new_element.text = text
        if attributes is not None:
            for key, value in attributes.items():
                new_element.set(key, value)
        
        return new_element
    
    def add_element_with_namespace(self, parent_xpath: str, tag: str, namespace_prefix: str, 
                                  namespace_uri: str, text: str = None, 
                                  attributes: dict = None) -> Optional[ElementBase]:
        """
        添加带命名空间声明的元素

        :param parent_xpath: 父元素XPath
        :param tag: 元素标签名
        :param namespace_prefix: 命名空间前缀
        :param namespace_uri: 命名空间URI
        :param text: 元素文本内容
        :param attributes: 元素属性字典
        :return: 新添加的元素
        """
        parent = self.find_element(parent_xpath)
        if parent is None:
            return None
        
        # 注册命名空间
        self.add_namespace(namespace_prefix, namespace_uri)
        
        # 创建带命名空间的标签
        full_tag = f"{{{namespace_uri}}}{tag}"
        new_element = etree.SubElement(parent, full_tag, nsmap={namespace_prefix: namespace_uri})
        
        if text is not None:
            new_element.text = text
        if attributes is not None:
            for key, value in attributes.items():
                new_element.set(key, value)
        
        return new_element
    
    def set_element_namespace(self, xpath: str, namespace_prefix: str, namespace_uri: str) -> bool:
        """
        为现有元素设置命名空间

        :param xpath: 元素XPath
        :param namespace_prefix: 命名空间前缀
        :param namespace_uri: 命名空间URI
        :return: 是否设置成功
        """
        elements = self.find_elements(xpath)
        if not elements:
            return False
        
        # 注册命名空间
        self.add_namespace(namespace_prefix, namespace_uri)
        
        # 为所有匹配的元素添加命名空间声明
        for element in elements:
            # 通过创建新元素并替换旧元素的方式来添加命名空间
            parent = element.getparent()
            if parent is not None:
                # 创建带命名空间的新元素
                new_element = etree.Element(element.tag, nsmap={namespace_prefix: namespace_uri})
                new_element.text = element.text
                new_element.tail = element.tail
                
                # 复制属性
                for key, value in element.attrib.items():
                    new_element.set(key, value)
                
                # 复制子元素
                for child in element:
                    new_element.append(child)
                
                # 替换旧元素
                parent.replace(element, new_element)
        
        return True
    
    def set_element_namespace_declaration(self, xpath: str, prefix: str, uri: str) -> bool:
        """
        为元素设置命名空间声明（xmlns:prefix="uri"）

        :param xpath: 元素XPath
        :param prefix: 命名空间前缀
        :param uri: 命名空间URI
        :return: 是否设置成功
        """
        element = self.find_element(xpath)
        if element is None:
            return False
        
        # 创建新的命名空间映射
        nsmap = element.nsmap.copy() if element.nsmap else {}
        
        # 设置命名空间（None表示默认命名空间）
        nsmap[prefix if prefix else None] = uri
        
        # 重新创建元素以应用新的命名空间映射
        parent = element.getparent()
        if parent is not None:
            # 创建新元素并复制所有内容
            new_element = etree.Element(element.tag, nsmap=nsmap)
            new_element.text = element.text
            new_element.tail = element.tail
            
            # 复制属性
            for key, value in element.attrib.items():
                new_element.set(key, value)
            
            # 复制子元素
            for child in element:
                new_element.append(child)
            
            # 替换旧元素
            parent.replace(element, new_element)
        else:
            # 如果是根元素，直接更新
            element.clear()
            element.tag = element.tag  # 保持原标签
            element.nsmap = nsmap
        
        # 更新内部命名空间映射
        self.add_namespace(prefix, uri)
        
        return True
    
    def get_element_namespaces(self, xpath: str) -> Dict[str, str]:
        """
        获取元素的命名空间声明

        :param xpath: 元素XPath
        :return: 命名空间字典
        """
        element = self.find_element(xpath)
        if element is not None:
            return element.nsmap if hasattr(element, 'nsmap') else {}
        return {}
    
    def add_attribute_with_namespace(self, xpath: str, attr_name: str, attr_value: str, 
                                   namespace_prefix: str = None, namespace_uri: str = None) -> bool:
        """
        为元素添加带命名空间的属性

        :param xpath: 元素XPath
        :param attr_name: 属性名
        :param attr_value: 属性值
        :param namespace_prefix: 属性命名空间前缀
        :param namespace_uri: 属性命名空间URI
        :return: 是否添加成功
        """
        element = self.find_element(xpath)
        if element is None:
            return False
        
        # 构造带命名空间的属性名
        if namespace_prefix and namespace_uri:
            # 注册命名空间
            self.add_namespace(namespace_prefix, namespace_uri)
            full_attr_name = f"{{{namespace_uri}}}{attr_name}"
            element.set(full_attr_name, attr_value)
        else:
            element.set(attr_name, attr_value)
        
        return True